import uuid

from pymilvus import Collection, FieldSchema, DataType, CollectionSchema

from bot.openai_bot import OpenAIBot
from conf.config import logger
from db.milvus.text_chunk_dao import BaseTextChunkDAO


class BuildingInfoTextChunkDAO(BaseTextChunkDAO):
    """ 楼盘信息文本块数据访问对象 """

    def __init__(self, *args, **kwargs):
        super(BuildingInfoTextChunkDAO, self).__init__(*args, **kwargs)
        self.alias = str(uuid.uuid4())  # 给本次的milvus链接起个名字，后续释放链接要用到该名字
        self.collection_name = "text_chunk_of_building_info"

    def create_collection(self, dim: int):
        """
        创建集合
        :param dim: 矢量维度
        :return:
        """
        milvus_client = self.create_milvus_conn()
        try:

            fields = [
                FieldSchema(name="text_chunk_id", dtype=DataType.INT64, is_primary=True, auto_id=False, description="文本块id"),
                FieldSchema(name="text_chunk_vector", dtype=DataType.FLOAT_VECTOR, dim=dim, description="文本块对应的向量"),
            ]
            schema = CollectionSchema(fields=fields, description='楼盘信息文本块')
            collection = Collection(name=self.collection_name, schema=schema, using=self.alias)

            index_params = {
                'metric_type': "IP",
                'index_type': "FLAT",
                'params': {}
            }
            collection.create_index(field_name='text_chunk_vector', index_params=index_params)
        finally:
            if "milvus_client" in dir():
                milvus_client.disconnect()


def main():
    # # 步骤1：创建集合
    # building_info_text_chunk_dao = BuildingInfoTextChunkDAO()
    # building_info_text_chunk_dao.create_collection(dim=1536)
    #
    # # 步骤2：数据向量化
    # building_info_text_chunk_dao = BuildingInfoTextChunkDAO()
    # building_info_text_chunk_dao.data_to_embeddings(csv_filename="building_info_kejizhiguang.csv")

    # 步骤3：查找向量
    openai_bot = OpenAIBot()
    vector = openai_bot.embeddings(text_list=["保利天汇的咨询电话是多少？", ])[0]
    building_info_text_chunk_dao = BuildingInfoTextChunkDAO()
    res = building_info_text_chunk_dao.get_ids(vector=vector, top_k=5)
    logger.info(f"{res}")


if __name__ == '__main__':
    main()
